package jmeter.javarequest;

import com.google.gson.Gson;
import org.apache.jmeter.protocol.java.sampler.AbstractJavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;

import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.Properties;

/**
 * flink压力测试，第一步，装载压测数据到kafka的flink-tps-data topic
 *
 * @author evoup
 */
public class JavaRequestKafkaDataGenerator extends AbstractJavaSamplerClient {
    private Producer<String, String> producer;
    private Long counter = 0L;
    private static final String ENCODING = "UTF-8";

    @Override
    public void setupTest(JavaSamplerContext context) {
        Properties props = new Properties();
        props.put("bootstrap.servers", "172.16.30.217:9092");
        props.put("acks", "all");
        props.put("retries", 0);
        props.put("batch.size", 16384);
        props.put("linger.ms", 1);
        props.put("buffer.memory", 33554432);
        props.put("key.serializer", "org.apache.kafka.common.serialization.StringSerializer");
        props.put("value.serializer", "org.apache.kafka.common.serialization.StringSerializer");
        producer = new KafkaProducer<>(props);
    }

    @Override
    public SampleResult runTest(JavaSamplerContext javaSamplerContext) {
        SampleResult result = new SampleResult();
        InputJson inputJson = new InputJson();
        inputJson.setId(counter++);
        inputJson.setEventTime(System.currentTimeMillis());
        //20B: go into flink messag
        //50B: go into flink message,go into flink message,go int
        String msg = "go into flink messag";
        inputJson.setMsg(msg);
        sampleResultStart(result, msg);
        try {
            producer.send(new ProducerRecord<String, String>("flink-tps-data" +
                    "" +
                    "", null, new Gson().toJson(inputJson)));
            sampleResultSuccess(result, null);
        } catch (Exception e) {
            sampleResultFailed(result, "500", e);
        }
        return result;
    }

    @Override
    public void teardownTest(JavaSamplerContext context) {
        producer.close();
        super.teardownTest(context);
    }

    private void sampleResultStart(SampleResult result, String data) {
        result.setSamplerData(data);
        result.sampleStart();
    }

    private void sampleResultSuccess(SampleResult result, /* @Nullable */ String response) {
        result.sampleEnd();
        result.setSuccessful(true);
        result.setResponseCodeOK();
        if (response != null) {
            result.setResponseData(response, "");
        } else {
            result.setResponseData("No response required", ENCODING);
        }
    }

    private void sampleResultFailed(SampleResult result, String reason) {
        result.sampleEnd();
        result.setSuccessful(false);
        result.setResponseCode(reason);
    }

    private void sampleResultFailed(SampleResult result, String reason, Exception exception) {
        sampleResultFailed(result, reason);
        result.setResponseMessage("Exception: " + exception);
        result.setResponseData(getStackTrace(exception), ENCODING);
    }

    private String getStackTrace(Exception exception) {
        StringWriter stringWriter = new StringWriter();
        exception.printStackTrace(new PrintWriter(stringWriter));
        return stringWriter.toString();
    }

    private class InputJson {
        Long id;
        Long eventTime;
        String msg;

        public Long getId() {
            return id;
        }

        public void setId(Long id) {
            this.id = id;
        }

        public Long getEventTime() {
            return eventTime;
        }

        public void setEventTime(Long eventTime) {
            this.eventTime = eventTime;
        }

        public String getMsg() {
            return msg;
        }

        public void setMsg(String msg) {
            this.msg = msg;
        }
    }

}
